# ---------------------------------------------------------------------------------------------------
# CLIP-DINOiser
# authors: Monika Wysoczanska, Warsaw University of Technology & Oriane Simeoni, valeo.ai
# ---------------------------------------------------------------------------------------------------

from typing import Optional

import cv2
import numpy as np
from mmengine.registry import init_default_scope
from mmseg.datasets.builder import PIPELINES
from mmseg.datasets.pipelines import ImageToTensor, to_tensor

init_default_scope('mmseg')


@PIPELINES.register_module(force=True)
class ToRGB:
    def __call__(self, results):
        return self.transform(results)

    def transform(self, results: dict) -> Optional[dict]:
        
        results['img'] = cv2.cvtColor(results['img'], cv2.COLOR_BGR2RGB)
        return results


# ________________
# Modified version from mmcv to directly convert a tensor to float
# MAKE SURE YOU USE IT ONLY FOR IMAGES


@PIPELINES.register_module(force=True)
class ImageToTensorV2(ImageToTensor):


    def __init__(self, keys: dict) -> None:
        super(ImageToTensorV2, self).__init__(keys)
        self.keys = keys

    def __call__(self, results):
        return self.transform(results)

    def transform(self, results: dict) -> dict:
        
        for key in self.keys:
            img = results[key]
            if len(img.shape) < 3:
                img = np.expand_dims(img, -1)

            results[key] = (to_tensor(img.copy()).permute(2, 0, 1)).contiguous() / 255.

        return results
